use anyhow::Result;
use rig::client::EmbeddingsClient;
use rig::vector_store::request::VectorSearchRequest;
use rig::{
    Embed, embeddings::EmbeddingsBuilder, providers::openai::{Client, TEXT_EMBEDDING_ADA_002},
    vector_store::VectorStoreIndex,
};
use rig_sqlite::{Column, ColumnValue, SqliteVectorStore, SqliteVectorStoreTable};
use rusqlite::ffi::{sqlite3, sqlite3_api_routines, sqlite3_auto_extension};
use serde::{Deserialize, Serialize};
use sqlite_vec::sqlite3_vec_init;
use std::env;
use tokio_rusqlite::Connection;
use uuid::Uuid;

use crate::database::KnowledgeItem;
use rig_app::agents::AgentManager;
use rig_app::agent_types::{AgentCategory, EnhancedAgentConfig};

#[derive(Embed, Clone, Debug, Deserialize, Serialize)]
pub struct KnowledgeDocument {
    pub id: String,
    pub knowledge_item_id: i64,
    #[embed]
    pub content: String,
    pub title: String,
    pub item_type: String,
    pub created_at: String,
}

impl SqliteVectorStoreTable for KnowledgeDocument {
    fn name() -> &'static str {
        "knowledge_documents"
    }

    fn schema() -> Vec<Column> {
        vec![
            Column::new("id", "TEXT PRIMARY KEY"),
            Column::new("knowledge_item_id", "TEXT"),
            Column::new("content", "TEXT"),
            Column::new("title", "TEXT"),
            Column::new("item_type", "TEXT"),
            Column::new("created_at", "TEXT"),
        ]
    }

    fn id(&self) -> String {
        self.id.clone()
    }

    fn column_values(&self) -> Vec<(&'static str, Box<dyn ColumnValue>)> {
        vec![
            ("id", Box::new(self.id.clone())),
            ("knowledge_item_id", Box::new(self.knowledge_item_id.to_string())),
            ("content", Box::new(self.content.clone())),
            ("title", Box::new(self.title.clone())),
            ("item_type", Box::new(self.item_type.clone())),
            ("created_at", Box::new(self.created_at.clone())),
        ]
    }
}

type SqliteExtensionFn =
    unsafe extern "C" fn(*mut sqlite3, *mut *mut i8, *const sqlite3_api_routines) -> i32;

pub struct VectorSearchManager {
    openai_client: Client,
    vector_store: SqliteVectorStore<rig::providers::openai::EmbeddingModel, KnowledgeDocument>,
    agent_manager: Option<AgentManager>,
}

// 手动实现 Debug trait，因为某些字段可能不支持 Debug
impl std::fmt::Debug for VectorSearchManager {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("VectorSearchManager")
            .field("openai_client", &"<OpenAI Client>")
            .field("vector_store", &"<SQLite Vector Store>")
            .field("agent_manager", &"<Agent Manager>")
            .finish()
    }
}

impl VectorSearchManager {
    pub async fn new(db_path: &str) -> Result<Self> {
        // Initialize OpenAI client
        let openai_api_key = env::var("OPENAI_API_KEY")
            .map_err(|_| anyhow::anyhow!("OPENAI_API_KEY environment variable not set"))?;
        let openai_client = Client::new(&openai_api_key);

        // Initialize the sqlite-vec extension
        unsafe {
            sqlite3_auto_extension(Some(std::mem::transmute::<*const (), SqliteExtensionFn>(
                sqlite3_vec_init as *const (),
            )));
        }

        // Initialize SQLite connection
        let conn = Connection::open(db_path).await?;

        // Select the embedding model
        let model = openai_client.embedding_model(TEXT_EMBEDDING_ADA_002);

        // Initialize SQLite vector store
        let vector_store = SqliteVectorStore::new(conn, &model).await?;

        Ok(Self {
            openai_client,
            vector_store,
            agent_manager: None,
        })
    }

    pub async fn new_with_agent_manager(db_path: &str, agent_manager: AgentManager) -> Result<Self> {
        // Initialize OpenAI client
        let openai_api_key = env::var("OPENAI_API_KEY")
            .map_err(|_| anyhow::anyhow!("OPENAI_API_KEY environment variable not set"))?;
        let openai_client = Client::new(&openai_api_key);

        // Initialize the sqlite-vec extension
        unsafe {
            sqlite3_auto_extension(Some(std::mem::transmute::<*const (), SqliteExtensionFn>(
                sqlite3_vec_init as *const (),
            )));
        }

        // Initialize SQLite connection
        let conn = Connection::open(db_path).await?;

        // Select the embedding model
        let model = openai_client.embedding_model(TEXT_EMBEDDING_ADA_002);

        // Initialize SQLite vector store
        let vector_store = SqliteVectorStore::new(conn, &model).await?;

        Ok(Self {
            openai_client,
            vector_store,
            agent_manager: Some(agent_manager),
        })
    }

    /// Get embedding model based on agent configuration
    async fn get_embedding_model_for_agent(&self, agent_id: Option<i64>) -> Result<rig::providers::openai::EmbeddingModel> {
        if let (Some(agent_manager), Some(agent_id)) = (&self.agent_manager, agent_id) {
            // Try to get agent configuration
            if let Ok(Some(agent_config)) = agent_manager.get_agent(agent_id).await {
                // Check if this is an embed type agent
                // The agent_config is already an EnhancedAgentConfig, no need to deserialize
                if agent_config.category == AgentCategory::Embed {
                    // Use the agent's embedding configuration
                    // For now, we'll use the default model, but this can be extended
                    // to support different embedding models based on agent config
                    return Ok(self.openai_client.embedding_model(TEXT_EMBEDDING_ADA_002));
                }
            }
        }
        
        // Fallback to default embedding model
        Ok(self.openai_client.embedding_model(TEXT_EMBEDDING_ADA_002))
    }

    pub async fn add_knowledge_item(&self, item: &KnowledgeItem) -> Result<()> {
        self.add_knowledge_item_with_agent(item, None).await
    }

    pub async fn add_knowledge_item_with_agent(&self, item: &KnowledgeItem, agent_id: Option<i64>) -> Result<()> {
        if let Some(content) = &item.content {
            if !content.trim().is_empty() {
                let document = KnowledgeDocument {
                    id: Uuid::new_v4().to_string(),
                    knowledge_item_id: item.id,
                    content: content.clone(),
                    title: item.title.clone(),
                    item_type: item.item_type.clone(),
                    created_at: item.created_at.clone(),
                };

                let model = self.get_embedding_model_for_agent(agent_id).await?;
                let embeddings = EmbeddingsBuilder::new(model)
                    .documents(vec![document])?
                    .build()
                    .await?;

                self.vector_store.add_rows(embeddings).await?;
            }
        }
        Ok(())
    }

    pub async fn search_similar(
        &self,
        query: &str,
        limit: usize,
    ) -> Result<Vec<(f64, KnowledgeDocument)>> {
        self.search_similar_with_agent(query, limit, None).await
    }

    pub async fn search_similar_with_agent(
        &self,
        query: &str,
        limit: usize,
        agent_id: Option<i64>,
    ) -> Result<Vec<(f64, KnowledgeDocument)>> {
        let model = self.get_embedding_model_for_agent(agent_id).await?;
        let index = self.vector_store.clone().index(model);
        
        let req = VectorSearchRequest::builder()
            .samples(limit as u64)
            .query(query)
            .build()?;

        let results = index
            .top_n::<KnowledgeDocument>(req)
            .await?
            .into_iter()
            .map(|(score, _id, doc)| (score, doc))
            .collect::<Vec<_>>();

        Ok(results)
    }

    pub async fn search_similar_ids(&self, query: &str, limit: usize) -> Result<Vec<(f64, String)>> {
        self.search_similar_ids_with_agent(query, limit, None).await
    }

    pub async fn search_similar_ids_with_agent(&self, query: &str, limit: usize, agent_id: Option<i64>) -> Result<Vec<(f64, String)>> {
        let model = self.get_embedding_model_for_agent(agent_id).await?;
        let index = self.vector_store.clone().index(model);
        
        let req = VectorSearchRequest::builder()
            .samples(limit as u64)
            .query(query)
            .build()?;

        let results = index.top_n_ids(req).await?.into_iter().collect::<Vec<_>>();

        Ok(results)
    }

    pub async fn update_knowledge_item(&self, item: &KnowledgeItem) -> Result<()> {
        self.update_knowledge_item_with_agent(item, None).await
    }

    pub async fn update_knowledge_item_with_agent(&self, item: &KnowledgeItem, agent_id: Option<i64>) -> Result<()> {
        // For updates, we might want to remove old embeddings and add new ones
        // This is a simplified approach - in production, you might want to track embeddings more carefully
        self.remove_knowledge_item(item.id).await?;
        self.add_knowledge_item_with_agent(item, agent_id).await?;
        Ok(())
    }

    pub async fn remove_knowledge_item(&self, _knowledge_item_id: i64) -> Result<()> {
        // Note: This is a simplified implementation
        // In a real implementation, you'd want to remove specific embeddings
        // For now, we'll rely on the database foreign key constraints
        tracing::warn!("Vector removal not fully implemented - relying on DB constraints");
        Ok(())
    }

    pub async fn batch_add_knowledge_items(&self, items: Vec<&KnowledgeItem>) -> Result<()> {
        self.batch_add_knowledge_items_with_agent(items, None).await
    }

    pub async fn batch_add_knowledge_items_with_agent(&self, items: Vec<&KnowledgeItem>, agent_id: Option<i64>) -> Result<()> {
        let documents: Vec<KnowledgeDocument> = items
            .into_iter()
            .filter_map(|item| {
                if let Some(content) = &item.content {
                    if !content.trim().is_empty() {
                        Some(KnowledgeDocument {
                            id: Uuid::new_v4().to_string(),
                            knowledge_item_id: item.id,
                            content: content.clone(),
                            title: item.title.clone(),
                            item_type: item.item_type.clone(),
                            created_at: item.created_at.clone(),
                        })
                    } else {
                        None
                    }
                } else {
                    None
                }
            })
            .collect();

        if !documents.is_empty() {
            let model = self.get_embedding_model_for_agent(agent_id).await?;
            let embeddings = EmbeddingsBuilder::new(model)
                .documents(documents)?
                .build()
                .await?;

            self.vector_store.add_rows(embeddings).await?;
        }

        Ok(())
    }
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct SearchResult {
    pub knowledge_item: KnowledgeItem,
    pub similarity_score: f32,
    pub search_type: String, // "vector" or "text"
}

#[derive(Debug)]
pub struct HybridSearchManager {
    vector_search: VectorSearchManager,
}

impl HybridSearchManager {
    pub async fn new(db_path: &str) -> Result<Self> {
        let vector_search = VectorSearchManager::new(db_path).await?;
        Ok(Self { vector_search })
    }

    pub async fn new_with_agent_manager(db_path: &str, agent_manager: AgentManager) -> Result<Self> {
        let vector_search = VectorSearchManager::new_with_agent_manager(db_path, agent_manager).await?;
        Ok(Self { vector_search })
    }

    pub async fn hybrid_search(
        &self,
        query: &str,
        limit: usize,
        _vector_weight: f32,
    ) -> Result<Vec<SearchResult>> {
        self.hybrid_search_with_agent(query, limit, _vector_weight, None).await
    }

    pub async fn hybrid_search_with_agent(
        &self,
        query: &str,
        limit: usize,
        _vector_weight: f32,
        agent_id: Option<i64>,
    ) -> Result<Vec<SearchResult>> {
        // Perform vector search with agent configuration
        let vector_results = self.vector_search.search_similar_with_agent(query, limit, agent_id).await?;

        // Convert to SearchResult format
        let mut results: Vec<SearchResult> = vector_results
            .into_iter()
            .map(|(score, doc)| {
                let created_at = doc.created_at.clone();
                SearchResult {
                    knowledge_item: KnowledgeItem {
                        id: doc.knowledge_item_id,
                        uuid: Uuid::new_v4().to_string(), // This should be fetched from DB
                        title: doc.title,
                        content: Some(doc.content),
                        item_type: doc.item_type,
                        url: None,
                        file_path: None,
                        file_size: None,
                        mime_type: None,
                        created_at: created_at.clone(),
                        updated_at: created_at,
                        tags: None,
                        metadata: None,
                    },
                    similarity_score: score as f32,
                    search_type: "vector".to_string(),
                }
            })
            .collect();

        // Sort by similarity score (descending)
        results.sort_by(|a, b| b.similarity_score.partial_cmp(&a.similarity_score).unwrap());

        // Limit results
        results.truncate(limit);

        Ok(results)
    }

    pub async fn add_knowledge_item(&self, item: &KnowledgeItem) -> Result<()> {
        self.vector_search.add_knowledge_item(item).await
    }

    pub async fn add_knowledge_item_with_agent(&self, item: &KnowledgeItem, agent_id: Option<i64>) -> Result<()> {
        self.vector_search.add_knowledge_item_with_agent(item, agent_id).await
    }

    pub async fn update_knowledge_item(&self, item: &KnowledgeItem) -> Result<()> {
        self.vector_search.update_knowledge_item(item).await
    }

    pub async fn update_knowledge_item_with_agent(&self, item: &KnowledgeItem, agent_id: Option<i64>) -> Result<()> {
        self.vector_search.update_knowledge_item_with_agent(item, agent_id).await
    }

    pub async fn remove_knowledge_item(&self, knowledge_item_id: i64) -> Result<()> {
        self.vector_search.remove_knowledge_item(knowledge_item_id).await
    }

    pub async fn batch_add_knowledge_items(&self, items: Vec<&KnowledgeItem>) -> Result<()> {
        self.vector_search.batch_add_knowledge_items(items).await
    }

    pub async fn batch_add_knowledge_items_with_agent(&self, items: Vec<&KnowledgeItem>, agent_id: Option<i64>) -> Result<()> {
        self.vector_search.batch_add_knowledge_items_with_agent(items, agent_id).await
    }
}
